// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once
#include <memory>
#include <unordered_set>
#include <vector>
#include <pollux/core/plan_node.h>
#include <melon/container/f14_map.h>

namespace kumo::pollux::core {
    class QueryConfig;

    /// Gives hints on how to execute the fragment of a plan.
    enum class ExecutionStrategy {
        /// Process splits as they come in any available driver.
        kUngrouped,
        /// Process splits from each split group only in one driver.
        /// It is used when split groups represent separate partitions of the data on
        /// the grouping keys or join keys. In that case it is sufficient to keep only
        /// the keys from a single split group in a hash table used by group-by or
        /// join.
        kGrouped,
    };

    std::string executionStrategyToString(ExecutionStrategy strategy);

    /// Contains some information on how to execute the fragment of a plan.
    /// Used to construct Task.
    struct PlanFragment {
        /// Top level (root) Plan Node.
        std::shared_ptr<const core::PlanNode> planNode;
        ExecutionStrategy executionStrategy{ExecutionStrategy::kUngrouped};
        int numSplitGroups{0};

        /// Contains leaf plan nodes that need to be executed in the grouped mode.
        std::unordered_set<PlanNodeId> groupedExecutionLeafNodeIds;

        /// Returns true if the fragment uses grouped execution strategy meaning that
        /// at least one pipeline has a leaf node that should run grouped execution.
        /// Note that it does not mean that all pipelines run grouped execution -
        /// some leaf nodes might still run ungrouped execution.
        inline bool isGroupedExecution() const {
            return executionStrategy == ExecutionStrategy::kGrouped;
        }

        /// Returns true for leaf nodes that use grouped execution, false otherwise.
        inline bool leafNodeRunsGroupedExecution(const PlanNodeId &planNodeId) const {
            return groupedExecutionLeafNodeIds.find(planNodeId) !=
                   groupedExecutionLeafNodeIds.end();
        }

        PlanFragment() = default;

        explicit PlanFragment(std::shared_ptr<const core::PlanNode> topNode)
            : planNode(std::move(topNode)) {
        }

        PlanFragment(
            std::shared_ptr<const core::PlanNode> topNode,
            ExecutionStrategy strategy,
            int numberOfSplitGroups,
            const std::unordered_set<PlanNodeId> &groupedExecLeafNodeIds)
            : planNode(std::move(topNode)),
              executionStrategy(strategy),
              numSplitGroups(numberOfSplitGroups),
              groupedExecutionLeafNodeIds(groupedExecLeafNodeIds) {
        }

        /// Returns true if the spilling is enabled and there is at least one node in
        /// the plan, whose operator can spill. Returns false otherwise.
        bool canSpill(const QueryConfig &queryConfig) const;
    };
} // namespace kumo::pollux::core

template<>
struct fmt::formatter<kumo::pollux::core::ExecutionStrategy>
        : formatter<int> {
    auto format(
        const kumo::pollux::core::ExecutionStrategy &s,
        format_context &ctx) const {
        return formatter<int>::format(static_cast<int>(s), ctx);
    }
};
